import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


def load_model_and_tokenizer(checkpoint_path):
    # 根据设备配置加载模型
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    tokenizer = AutoTokenizer.from_pretrained(
        checkpoint_path,
        resume_download=True,
    )
    tokenizer.pad_token = tokenizer.eos_token  # 设置pad token
    model = AutoModelForCausalLM.from_pretrained(
        checkpoint_path,
        torch_dtype=torch.float32,
        device_map=device,
        resume_download=True,
        trust_remote_code=True
    )
    return model,tokenizer


def chat_with_llm(model, tokenizer, user_query: str,  max_new_tokens: int = 128):
    """
    实现对话功能（支持历史对话管理）

    Args:
        model: 加载后的AutoModelForCausalLM模型
        tokenizer: 对应的Tokenizer
        user_query: 当前用户输入
        max_new_tokens: 生成回答的最大长度（避免生成过长）
    Returns:
        bot_response: 模型生成的回答
        new_history: 更新后的对话历史
    """

    # 1. 构建对话prompt（按模型要求的格式拼接历史+当前查询）
    # 不同模型的prompt格式不同！需参考模型官方文档（以下为Qwen/Baichuan通用对话格式）
    prompt = ""
    # 添加当前用户查询
    prompt += f"用户：{user_query}\n助手："

    # 2. 文本→Token序列（模型输入格式）
    inputs = tokenizer(
        prompt,
        return_tensors="pt",  # 返回PyTorch张量
        padding=True,
        truncation=True,  # 截断过长的prompt（避免超出模型最大上下文长度）
        max_length=tokenizer.model_max_length - max_new_tokens  # 预留生成空间
    ).to(model.device)  # 将输入移到模型所在设备（GPU/CPU）

    # 3. 模型生成回答（核心步骤）
    with torch.no_grad():  # 禁用梯度计算，节省显存+提升速度
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,  # 最大生成长度
            do_sample=True,  # 采样生成（避免 deterministic 重复）
            temperature=0.7,  # 采样温度（越高越随机，0.7为平衡值）
            top_p=0.9,  # 核采样（只从概率前90%的token中选）
            repetition_penalty=1.1,  # 重复惩罚（减少回答重复，1.1为常用值）
            eos_token_id=tokenizer.eos_token_id,  # 生成终止符（遇到eos停止）
            pad_token_id=tokenizer.pad_token_id  # 处理padding的token
        )

    # 4. Token序列→文本（解码回答）
    # 只取“生成的部分”（排除原始prompt的token）
    bot_response = tokenizer.decode(
        outputs[0][len(inputs["input_ids"][0]):],  # 从prompt结束位置开始解码
        skip_special_tokens=True  # 跳过特殊token（如eos、pad）
    ).strip()

    return bot_response

if __name__ == '__main__':

    model_name='AiLab-0.1B'
    model_path=f"../data/{model_name}"
    print(f"正在加载模型：{model_path}...")
    model, tokenizer=load_model_and_tokenizer(model_path)
    print(model)
    print("模型加载完成！")

    # 2. 多轮对话交互
    history = None
    print("\n开始对话（输入'退出'结束）：")
    while True:
        user_input = input("你：")
        if user_input.lower() in ["退出", "quit", "exit"]:
            print("助手：再见！")
            break
        # 调用对话函数
        response = chat_with_llm(model, tokenizer, user_input)
        print(f"助手：{response}")